from ultralytics import YOLO  # 导入Yolo模块
import os
import  torch

# 关闭gpu 训练
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # This will make PyTorch only use the CPU
# Now we can create tensors on CPU even if we have a GPU available.
x=torch.tensor([1,0])
print(x) # This will output: tensor([1.0])

# 加载yolov11的预训练模型，这个模型是yolov8魔改，
# 我们将它作为基础模型，在该模型的基础上，训练安全帽模型
model = YOLO("model/yolov8n.pt")


if __name__ == '__main__':
    # 训练用户自定义的数据集，数据的配置保存在 safehat.yaml 中，epochs等于100表示100轮迭代
    # model.train(data="sagehat.yml", epochs=100, imgsz=640, batch=16, device='0')
    model.train(data="sagehat.yml", epochs=300, imgsz=640, batch=16)

    # 使用验证集验证效果
    model.val()

# yolo task=detect mode=train model=yolov8n.pt data=sagehat.yml epochs=300 imgsz=640 batch=16
# yolo task=detect mode=train model=last.pt data=sagehat.yml epochs=100 imgsz=640 batch=16